#pragma once

#include <memory>
#include <algorithm>
#include <iostream>

#include "contrib/Eigen/Dense"
#include "contrib/Eigen/Eigenvalues"

#pragma GCC optimize ("unroll-loops")

class Image {
    std::unique_ptr<float[]> BaseData;

public:
    Image();
    Image(size_t width, size_t height);
    Image(float* ptr, size_t width, size_t height);
    Image(Image&& image);

    Image& operator=(Image&& image);
    Image& operator+=(const Image& image);
    Image& operator/=(float v);

    Image DeepCopy() const;
    static Image Zeros(size_t width, size_t height);
    static Image Ones(size_t width, size_t height);

    bool IsEmpty() const;

    size_t Width, Height;
    float* Data;
};

Image SqrSobelH(const Image& input);
Image SqrSobelV(const Image& input);

void ToIntegral(Image& image);

template<int N>
Image BoxFilterByIntegralWithPadding(const Image& input) {
    Image result = Image(input.Width - N, input.Height - N); // Ignore last row and column
    float* resPtr = result.Data;
    for (size_t i = 0; i < result.Height; ++i) {
        float* inTopLeft = input.Data + input.Width * i;
        float* inTopRight = input.Data + input.Width * i + N;
        float* inBotLeft = input.Data + input.Width * (i + N);
        float* inBotRight = input.Data + input.Width * (i + N) + N;
        for (size_t j = 0; j < result.Width; ++j, ++resPtr, ++inTopLeft, ++inTopRight, ++inBotLeft, ++inBotRight) {
            *resPtr = *inBotRight + *inTopLeft - *inTopRight - *inBotLeft;
        }
    }
    return result;
}

template<int N>
Image EdgesMeasure(const Image& image) {
    Image h = SqrSobelH(image);
    ToIntegral(h);
    h = BoxFilterByIntegralWithPadding<N>(h);

    Image v = SqrSobelV(image);
    ToIntegral(v);
    v = BoxFilterByIntegralWithPadding<N>(v);

    h += v;
    return h;
}

template<int N>
void UpdateCovMatrix(const Image& input, Image& cov) {
    for (size_t i = 0, iCov = 0; i < N * N; ++i) {
        float t = input.Data[i];
        for (size_t j = 0; j < N * N; ++j, ++iCov) {
            cov.Data[iCov] += t * input.Data[j];
        }
    }
}

template<int N>
Image GetCovMatrix(const Image& image) {
    Image cov = Image::Zeros(N * N, N * N);
    Image temp(N, N);

    size_t strideI = image.Height / 128;
    size_t strideJ = image.Width / 128;
    size_t s = 0;
    for (size_t i = 1; i < (image.Height - 2) - N + 1; i += strideI) {
        for (size_t j = 1; j < (image.Width - 2) - N + 1; j += strideJ) {
            for (size_t ii = 0, iTemp = 0; ii < N; ++ii) {
                for (size_t jj = 0; jj < N; ++jj, ++iTemp) {
                    temp.Data[iTemp] = image.Data[(i + ii) * image.Width + j + jj];
                }
            }
            UpdateCovMatrix<N>(temp, cov);
            ++s;
        }
    }

    cov /= ((image.Height - 2) - N + 1) * ((image.Width - 2) - N + 1) / 4 - 1;
    return cov;
}

template<int N>
Image GetCovMatrixByEdgesThrash(const Image& image, const Image& edges, float thrash) {
    Image cov = Image::Zeros(N * N, N * N);
    Image temp(N, N);

    size_t strideI = image.Height / 128;
    size_t strideJ = image.Width / 128;
    size_t s = 0;
    for (size_t i = 1; i < (image.Height - 2) - N + 1; i += strideI) {
        for (size_t j = 1; j < (image.Width - 2) - N + 1; j += strideJ) {
            if (edges.Data[(i - 1) * edges.Width + (j - 1)] > thrash) {
                continue;
            }

            for (size_t ii = 0, iTemp = 0; ii < N; ++ii) {
                for (size_t jj = 0; jj < N; ++jj, ++iTemp) {
                    temp.Data[iTemp] = image.Data[(i + ii) * image.Width + j + jj];
                }
            }
            UpdateCovMatrix<N>(temp, cov);
            ++s;
        }
    }

    if (s < N * N) {
        return Image();
    }

    cov /= s;
    return cov;
}

template<int N>
float GetLastEigenvalue(const Image& matrix) {
    Eigen::Matrix<float, N * N, N * N> eigenMatrix;

    const float* srcPtr = matrix.Data;
    float* destPtr = eigenMatrix.data();
    for (size_t i = 0; i < N * N * N * N; ++i) {
        destPtr[i] = srcPtr[i];
    }
    auto eivals = eigenMatrix.eigenvalues().real();

    float res = 100000000000;
    for (size_t i = 0; i < N * N; ++i) {
        res = std::min(res, std::abs(eivals[i]));
    }
    return res;
}